#include "satlab.h"
#include <iostream>
#include <omp.h>
#include <chrono>
#include <cassert>
#include <cctype> // for isspace
#include <cstring>
#include <lzma.h>
#include <fstream>
#include <cstdint>
#include <set>
#include <map>
#include <filesystem> // Required for std::filesystem::create_directories

using namespace std;

typedef chrono::high_resolution_clock Clock;
typedef chrono::duration<double> Duration;

SATLAB::SATLAB(int image_w, int image_h)
    : varNum(0), claNum(0), posColor(cv::Vec3b(0, 0, 255)), negColor(cv::Vec3b(255, 0, 0)), image_w(static_cast<size_t>(image_w)), image_h(static_cast<size_t>(image_h)) {
    if (this->image_w <= 0 || this->image_w > 10000 || this->image_h <= 0 || this->image_h > 10000) {
        throw std::runtime_error("Unreasonable image size: " + std::to_string(this->image_w) + "x" + std::to_string(this->image_h));
    }
    scope.reserve(this->image_h);
    for (size_t i = 0; i < this->image_h; i++) {
        scope.push_back(vector<Pixel>(this->image_w));
    }
    claSeg.reserve(this->image_h);
    for (size_t i = 0; i < this->image_h + 1; i++) {
        claSeg.push_back({0, 0});
    }
}

void SATLAB::allocate(int varNum, int claNum) {
    this->varNum = static_cast<size_t>(varNum);
    this->claNum = static_cast<size_t>(claNum);
    mat.init(claNum);
    
    // Calculate block sizes
    blockVarSize = (this->varNum + image_w - 1) / image_w;
    blockClauseSize = (this->claNum + image_h - 1) / image_h;
    
    // Initialize first block row
    currentBlockRowIndex = 0;
    initializeBlockRow(0);
}

void SATLAB::initializeBlockRow(size_t rowIndex) {
    size_t numBlockCols = (varNum + blockVarSize - 1) / blockVarSize;
    
    currentBlockRow.clear();
    currentBlockRow.reserve(numBlockCols);
    
    for (size_t blockCol = 0; blockCol < numBlockCols; blockCol++) {
        Block block;
        block.clauseStart = rowIndex * blockClauseSize;
        block.clauseEnd = std::min((rowIndex + 1) * blockClauseSize, claNum);
        block.varStart = blockCol * blockVarSize;
        block.varEnd = std::min((blockCol + 1) * blockVarSize, varNum);
        
        // Calculate block dimensions
        size_t blockHeight = block.clauseEnd - block.clauseStart;
        size_t blockWidth = block.varEnd - block.varStart;
        
        // Initialize pixel grid for this block
        block.pixels.resize(blockHeight);
        for (auto& row : block.pixels) {
            row.resize(blockWidth);
        }
        
        currentBlockRow.push_back(block);
    }
}

void SATLAB::clearBlockRow() {
    currentBlockRow.clear();
}

void SATLAB::processBlockRowIfComplete(const std::string& outputDir, size_t currentClauseIndex) {
    // Check if we've completed the current block row
    size_t nextBlockRowStart = (currentBlockRowIndex + 1) * blockClauseSize;
    
    if (currentClauseIndex >= nextBlockRowStart || currentClauseIndex >= claNum) {
        // Generate images for the current block row
        if (!currentBlockRow.empty()) {
            createCurrentBlockRowImages(outputDir, currentBlockRowIndex);
            
            // Clear the current block row to free memory
            clearBlockRow();
            
            // Initialize the next block row
            currentBlockRowIndex++;
            if (currentBlockRowIndex * blockClauseSize < claNum) {
                initializeBlockRow(currentBlockRowIndex);
            }
        }
    }
}

void SATLAB::updateAllStats() {
    updateAllStatsForBlockRow(currentBlockRowIndex);
}

void SATLAB::updateAllStatsForBlockRow(size_t targetBlockRow) {
    for (auto& row : scope) {
        for (auto& pixel : row) {
            pixel.pos = 0;
            pixel.neg = 0;
        }
    }
    
    // Clear current block row data
    for (auto& block : currentBlockRow) {
        for (auto& row : block.pixels) {
            for (auto& pixel : row) {
                pixel.pos = 0;
                pixel.neg = 0;
            }
        }
    }

    #pragma omp parallel for schedule(dynamic)
    for (uint64_t segDex = 0; segDex < (uint64_t)claSeg.size(); segDex++) {
        uint64_t head = (segDex == 0) ? 0 : claSeg[segDex - 1].dex;
        uint64_t tail = claSeg[segDex].dex;
        for (uint64_t clauseDex = head; clauseDex < tail; clauseDex++) {
            const auto& clause = mat[clauseDex];
            for (Lit lit : clause) {
                uint64_t varDex = std::abs(lit) - 1;
                uint64_t colDex = (static_cast<uint64_t>(image_w) * varDex) / varNum;
                if (segDex > (uint64_t)image_h) {
                    std::cerr << "Assertion failed: segDex > image_h. segDex = " << segDex << ", image_h = " << image_h << std::endl;
                    assert(false);
                }
                if (colDex > (uint64_t)image_w) {
                    std::cerr << "Assertion failed: colDex > image_w. colDex = " << colDex << ", image_w = " << image_w << std::endl;
                    assert(false);
                }
                if (scope.size() <= segDex) {
                    std::cerr << "Assertion failed: scope.size() <= segDex. scope.size() = " << scope.size() << ", segDex = " << segDex << std::endl;
                    assert(false);
                }
                if (lit > 0) {
                    #pragma omp atomic
                    scope[segDex][colDex].pos++;
                } else {
                    #pragma omp atomic
                    scope[segDex][colDex].neg++;
                }
                
                // Also update current block row data
                size_t blockRow = clauseDex / blockClauseSize;
                size_t blockCol = varDex / blockVarSize;
                
                // Only process if this clause belongs to the target block row
                if (blockRow == targetBlockRow && blockCol < currentBlockRow.size()) {
                    size_t localClauseDex = clauseDex - currentBlockRow[blockCol].clauseStart;
                    size_t localVarDex = varDex - currentBlockRow[blockCol].varStart;
                    
                    if (localClauseDex < currentBlockRow[blockCol].pixels.size() && 
                        localVarDex < currentBlockRow[blockCol].pixels[0].size()) {
                        if (lit > 0) {
                            #pragma omp atomic
                            currentBlockRow[blockCol].pixels[localClauseDex][localVarDex].pos++;
                        } else {
                            #pragma omp atomic
                            currentBlockRow[blockCol].pixels[localClauseDex][localVarDex].neg++;
                        }
                    }
                }
            }
        }
    }
}

cv::Mat& SATLAB::createImage(cv::Mat& image, double colorThreshRatio) {
    auto start = Clock::now();
    image = cv::Mat(image_h, image_w, CV_8UC3, cv::Scalar(0, 0, 0));
    for (size_t row = 0; row < image_h; row++) {
        for (size_t col = 0; col < image_w; col++) {
            const Pixel& stat = scope[row][col];
            int total = stat.pos + stat.neg;
            if (total > 0) {
                double R = stat.pos / (double)total;
                double B = stat.neg / (double)total;
                cv::Vec3b& pixel = image.at<cv::Vec3b>(row, col);
                int minColor = static_cast<int>(255 * colorThreshRatio);
                int colorRange = 255 - minColor;
                pixel[0] = static_cast<uchar>(minColor + colorRange * B);  // Blue
                pixel[2] = static_cast<uchar>(minColor + colorRange * R);  // Red
            }
        }
    }
    auto end = Clock::now();
    Duration scope_time = end - start;
    cout << "\U0001F551 Scope Image Creation: " << scope_time.count() << " seconds" << endl;
    return image;
}

cv::Mat& SATLAB::createBinScopeImage(cv::Mat& image, size_t binIndex, double colorThreshRatio) {
    auto start = Clock::now();
    
    // Validate bin index
    if (binIndex >= image_h) {
        throw std::runtime_error("Bin index " + std::to_string(binIndex) + " is out of range. Max: " + std::to_string(image_h - 1));
    }
    
    image = cv::Mat(1, image_w, CV_8UC3, cv::Scalar(0, 0, 0));
    
    // Get the specific bin's data
    const std::vector<Pixel>& binData = scope[binIndex];
    
    for (size_t col = 0; col < image_w; col++) {
        const Pixel& stat = binData[col];
        int total = stat.pos + stat.neg;
        if (total > 0) {
            double R = stat.pos / (double)total;
            double B = stat.neg / (double)total;
            cv::Vec3b& pixel = image.at<cv::Vec3b>(0, col);
            int minColor = static_cast<int>(255 * colorThreshRatio);
            int colorRange = 255 - minColor;
            pixel[0] = static_cast<uchar>(minColor + colorRange * B);  // Blue
            pixel[2] = static_cast<uchar>(minColor + colorRange * R);  // Red
        }
    }
    
    auto end = Clock::now();
    Duration scope_time = end - start;
    cout << "\U0001F551 Bin " << binIndex << " Scope Image Creation: " << scope_time.count() << " seconds" << endl;
    return image;
}

void SATLAB::createAllBinScopeImages(const std::string& outputDir, double colorThreshRatio) {
    auto start = Clock::now();
    
    // Create output directory if it doesn't exist
    std::filesystem::create_directories(outputDir);
    
    cv::Mat binImage;
    for (size_t binIndex = 0; binIndex < image_h; binIndex++) {
        createBinScopeImage(binImage, binIndex, colorThreshRatio);
        
        // Save the bin image
        std::string filename = outputDir + "/bin_" + std::to_string(binIndex) + ".png";
        cv::imwrite(filename, binImage);
    }
    
    auto end = Clock::now();
    Duration total_time = end - start;
    cout << "\U0001F551 All Bin Scope Images Created: " << total_time.count() << " seconds" << endl;
    cout << "\U0001F4C1 Saved " << image_h << " bin images to: " << outputDir << endl;
}

cv::Mat& SATLAB::createBlockImage(cv::Mat& image, size_t blockIndex, double colorThreshRatio) {
    auto start = Clock::now();
    
    // Validate block index
    if (blockIndex >= currentBlockRow.size()) {
        throw std::runtime_error("Block index " + std::to_string(blockIndex) + " is out of range. Max: " + std::to_string(currentBlockRow.size() - 1));
    }
    
    const Block& block = currentBlockRow[blockIndex];
    size_t height = block.pixels.size();
    size_t width = block.pixels[0].size();
    
    image = cv::Mat(height, width, CV_8UC3, cv::Scalar(0, 0, 0)); // Initialize with black
    
    for (size_t row = 0; row < height; row++) {
        for (size_t col = 0; col < width; col++) {
            const Pixel& stat = block.pixels[row][col];
            
            cv::Vec3b& pixel = image.at<cv::Vec3b>(row, col);
            
            if (stat.pos > 0 && stat.neg == 0) {
                // Only positive literals present - full green
                pixel[1] = 255; // Green channel
                pixel[0] = 0;    // Blue channel
                pixel[2] = 0;    // Red channel
            } else if (stat.neg > 0 && stat.pos == 0) {
                // Only negative literals present - full red
                pixel[2] = 255; // Red channel
                pixel[0] = 0;   // Blue channel
                pixel[1] = 0;   // Green channel
            } else if (stat.pos > 0 && stat.neg > 0) {
                // Both positive and negative literals present - mix of red and green
                pixel[2] = 255; // Red channel
                pixel[1] = 255; // Green channel
                pixel[0] = 0;   // Blue channel
            }
            // If both pos and neg are 0, pixel remains black (default)
        }
    }
    
    auto end = Clock::now();
    Duration scope_time = end - start;
    cout << "\U0001F551 Block " << blockIndex << " Image Creation: " << scope_time.count() << " seconds" << endl;
    return image;
}

void SATLAB::createCurrentBlockRowImages(const std::string& outputDir, size_t rowIndex, double colorThreshRatio) {
    auto start = Clock::now();
    
    // Create output directory if it doesn't exist
    std::filesystem::create_directories(outputDir);
    
    cv::Mat blockImage;
    size_t savedCount = 0;
    
    for (size_t blockIndex = 0; blockIndex < currentBlockRow.size(); blockIndex++) {
        // Count total literals in this block
        size_t totalLiterals = 0;
        const Block& block = currentBlockRow[blockIndex];
        
        for (const auto& row : block.pixels) {
            for (const auto& pixel : row) {
                totalLiterals += pixel.pos + pixel.neg;
            }
        }
        
        // Skip blocks with fewer than 100 literals
        if (totalLiterals < 100) {
            continue;
        }
        
        createBlockImage(blockImage, blockIndex, colorThreshRatio);
        
        // Save the block image with row index in filename
        std::string filename = outputDir + "/block_row" + std::to_string(rowIndex) + "_col" + std::to_string(blockIndex) + ".png";
        cv::imwrite(filename, blockImage);
        savedCount++;
    }
    
    auto end = Clock::now();
    Duration total_time = end - start;
    cout << "\U0001F551 Block Row " << rowIndex << " Images Created: " << total_time.count() << " seconds" << endl;
    cout << "\U0001F4C1 Saved " << savedCount << " block images (skipped " << (currentBlockRow.size() - savedCount) << " sparse blocks) to: " << outputDir << endl;
}

cv::Mat& SATLAB::createHoughTransformImage(cv::Mat& image, const cv::Mat& scopeImage) {
    cv::Mat gray;
    cv::cvtColor(scopeImage, gray, cv::COLOR_BGR2GRAY);
    std::vector<cv::Vec4i> lines;
    cv::HoughLinesP(gray, lines, 1, CV_PI / 180, 100, 100, 10);
    image = cv::Mat::zeros(gray.size(), CV_8UC3);
    for (const auto& line : lines) {
        cv::line(image, cv::Point(line[0], line[1]), cv::Point(line[2], line[3]), cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
    }
    return image;
}

cv::Mat& SATLAB::createFFTImage(cv::Mat& image, const cv::Mat& scopeImage) {
    cv::Mat gray;
    cv::cvtColor(scopeImage, gray, cv::COLOR_BGR2GRAY);
    cv::Mat floatImg;
    gray.convertTo(floatImg, CV_32F);
    cv::Mat padded;
    int m = cv::getOptimalDFTSize(floatImg.rows);
    int n = cv::getOptimalDFTSize(floatImg.cols);
    cv::copyMakeBorder(floatImg, padded, 0, m - floatImg.rows, 0, n - floatImg.cols, cv::BORDER_CONSTANT, cv::Scalar::all(0));
    cv::Mat planes[] = {padded, cv::Mat::zeros(padded.size(), CV_32F)};
    cv::Mat complexI;
    cv::merge(planes, 2, complexI);
    cv::dft(complexI, complexI);
    cv::split(complexI, planes);
    cv::magnitude(planes[0], planes[1], planes[0]);
    cv::Mat magI = planes[0];
    magI += cv::Scalar::all(1);
    cv::log(magI, magI);
    magI = magI(cv::Rect(0, 0, floatImg.cols, floatImg.rows));
    int cx = magI.cols / 2;
    int cy = magI.rows / 2;
    cv::Mat q0(magI, cv::Rect(0, 0, cx, cy));
    cv::Mat q1(magI, cv::Rect(cx, 0, cx, cy));
    cv::Mat q2(magI, cv::Rect(0, cy, cx, cy));
    cv::Mat q3(magI, cv::Rect(cx, cy, cx, cy));
    cv::Mat tmp;
    q0.copyTo(tmp); q3.copyTo(q0); tmp.copyTo(q3);
    q1.copyTo(tmp); q2.copyTo(q1); tmp.copyTo(q2);
    cv::normalize(magI, magI, 0, 255, cv::NORM_MINMAX);
    magI.convertTo(image, CV_8U);
    cv::cvtColor(image, image, cv::COLOR_GRAY2BGR);
    return image;
}

cv::Mat& SATLAB::createMorphImage(cv::Mat& image, const cv::Mat& scopeImage) {
    cv::Mat gray;
    cv::cvtColor(scopeImage, gray, cv::COLOR_BGR2GRAY);
    int morphSize = 15;
    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(morphSize, morphSize));
    cv::Mat morph;
    cv::morphologyEx(gray, morph, cv::MORPH_OPEN, kernel);
    cv::cvtColor(morph, image, cv::COLOR_GRAY2BGR);
    return image;
}

char* readWS(char* p) { 
    while (isspace(*p)) ++p;
    return p; 
}

char* readNL(char* p) {
    while (*p != '\n' && (*p++ != '\0' || (exit(1), 0)));
    return ++p;
}

char* readNum(char* p, int* i) {
    p = readWS(p);
    int sgn = (*p == '-') ? (++p, -1) : 1;
    *i = 0;
    while (isdigit(*p)) *i = *i * 10 + (*p++ - '0');
    *i *= sgn;
    return p;
}


void SATLAB::parseFromMemory(char* data) {
    // First pass to get header info
    char* p = data;
    int varNum = 0;
    int claNum = 0;
    bool headerFound = false;
    
    while (*p != '\0') {
        p = readWS(p);
        if (*p == 'p' && !memcmp(p, "p cnf", 5)) {
            p += 5;
            p = readNum(p, &varNum);
            p = readNum(p, &claNum);
            if (varNum <= 0 || varNum > 100000000 || claNum <= 0 || claNum > 1000000000) {
                throw std::runtime_error("Unreasonable varNum or claNum in CNF file: " + std::to_string(varNum) + ", " + std::to_string(claNum));
            }
            allocate(varNum, claNum);
            headerFound = true;
            break;
        }
        p = readNL(p);
    }

    if (!headerFound) {
        throw runtime_error("Missing 'p cnf' line in the input data");
    }

    cout << "🔍 VarNum: " << varNum << endl;
    cout << "🔍 ClaNum: " << claNum << endl;
    // Calculate the size of each segment
    size_t segmentSize = (claNum + image_h - 1) / image_h;

    // After parsing header and before stats, resize scope and claSeg to match image_w and image_h
    scope.clear();
    scope.reserve(image_h);
    for (size_t i = 0; i < image_h; i++) {
        scope.push_back(std::vector<Pixel>(image_w));
    }
    claSeg.clear();
    claSeg.reserve(image_h);
    for (size_t i = 0; i < image_h + 1; i++) {
        claSeg.push_back({0, 0});
    }

    vector<Lit> lits;
    size_t claDex = 0;
    char* scanP = p;
    while (*scanP != '\0') {
        scanP = readWS(scanP);
        switch (*scanP) {
        case '%':
        case '\0':
            goto scan_done;
        case 'c':
        case 'p':
            scanP = readNL(scanP);
            break;
        default:
            int lit;
            scanP = readNum(scanP, &lit);
            if (lit == 0) {
                size_t segDex = claDex / segmentSize;
                claSeg[segDex] = {static_cast<size_t>(scanP - data), claDex};
                claDex++;
            }
        }
    }
scan_done:
    // Ensure the last segment is correctly set if not already
    if (claDex > 0) {
        size_t lastSegDex = claDex / segmentSize;
        if (lastSegDex < image_h) {
            claSeg[lastSegDex] = {static_cast<size_t>(scanP - data), claDex};
        }
    }

    #pragma omp parallel for schedule(dynamic)
    for (size_t segDex = 0; segDex < claSeg.size(); segDex++) {
        size_t head = (segDex == 0) ? 0 : claSeg[segDex - 1].ptr;
        size_t tail = claSeg[segDex].ptr;
        size_t startClaDex = (segDex == 0) ? 0 : (claSeg[segDex - 1].dex + 1);
        size_t endClaDex = claSeg[segDex].dex;

        char* p = data + head;
        vector<Lit> lits;
        int lit;
        size_t currentClaDex = startClaDex;

        while (p < data + tail && *p != '\0' && currentClaDex <= endClaDex) {
            p = readWS(p);
            switch (*p) {
            case '%':
            case '\0':
                if (!lits.empty()) {
                    cerr << "Unexpected end of clause" << endl;
                }
                break;
            case 'c':
            case 'p':
                p = readNL(p);
                break;
            default:
                p = readNum(p, &lit);
                if (lit != 0) {
                    lits.push_back(lit);
                } else if (!lits.empty()) {
                    sort(lits.begin(), lits.end(), [](int a, int b) { return abs(a) < abs(b); });
                    if (currentClaDex < mat.size()) {
                        mat[currentClaDex] = lits;
                        currentClaDex++;
                    } else {
                        cerr << "Clause index out of bounds" << endl;
                    }
                    lits.clear();
                }
            }
        }
    }

    // Handle any remaining literals (if last clause wasn't properly closed)
    if (!lits.empty()) {
        sort(lits.begin(), lits.end(), [](int a, int b) { return abs(a) < abs(b); });
        if (claDex < mat.size()) {
            mat[claDex++] = lits;
        } else {
            cerr << "Clause index out of bounds for remaining literals" << endl;
        }
        lits.clear();
    }

    if (image_w <= 0 || image_w > 10000 || image_h <= 0 || image_h > 10000) {
        throw std::runtime_error("Unreasonable image size after parsing: " + std::to_string(image_w) + "x" + std::to_string(image_h));
    }
}

void SATLAB::parseCompressed(const char* filename) {
    auto start_decompress = chrono::high_resolution_clock::now();
    
    // Open the XZ file
    FILE *inFile = fopen(filename, "rb");
    if (!inFile) {
        throw runtime_error("Cannot open input file");
    }

    // Initialize decoder
    lzma_stream strm = LZMA_STREAM_INIT;
    lzma_ret ret = lzma_stream_decoder(&strm, UINT64_MAX, LZMA_CONCATENATED);
    if (ret != LZMA_OK) {
        fclose(inFile);
        throw runtime_error("Failed to initialize LZMA decoder");
    }

    // Prepare buffers with larger size for better throughput
    const size_t BUFFER_SIZE = 1024 * 1024; // 1MB buffers
    vector<uint8_t> inBuf(BUFFER_SIZE);
    vector<uint8_t> outBuf(BUFFER_SIZE);
    vector<char> decompressed;
    decompressed.reserve(BUFFER_SIZE * 10); // Reserve space to reduce reallocations

    strm.next_out = outBuf.data();
    strm.avail_out = outBuf.size();

    // Decompress
    while (true) {
        if (strm.avail_in == 0) {
            strm.next_in = inBuf.data();
            strm.avail_in = fread(inBuf.data(), 1, inBuf.size(), inFile);
        }

        ret = lzma_code(&strm, strm.avail_in == 0 ? LZMA_FINISH : LZMA_RUN);

        if (strm.avail_out == 0 || ret == LZMA_STREAM_END) {
            size_t write_size = outBuf.size() - strm.avail_out;
            decompressed.insert(decompressed.end(), 
                              outBuf.data(), 
                              outBuf.data() + write_size);
            strm.next_out = outBuf.data();
            strm.avail_out = outBuf.size();
        }

        if (ret == LZMA_STREAM_END) {
            break;
        }

        if (ret != LZMA_OK) {
            lzma_end(&strm);
            fclose(inFile);
            throw runtime_error("Error during decompression");
        }
    }

    // Clean up
    lzma_end(&strm);
    fclose(inFile);

    auto end_decompress = Clock::now();
    Duration decompress_time = end_decompress - start_decompress;
    cout << "🕐 Decomp: " << decompress_time.count() << " seconds" << endl;

    // Ensure null termination
    decompressed.push_back('\0');

    // Parse the decompressed data 
    auto start_parse = Clock::now();
    parseFromMemory(decompressed.data());
    auto end_parse = Clock::now();
    Duration parse_time = end_parse - start_parse;
    cout << "🕑 Parsing: " << parse_time.count() << " seconds" << endl;
}

void SATLAB::parsePlainText(const char* filename) {
    auto start_read = Clock::now();
    
    ifstream fin(filename);
    if (!fin) {
        throw runtime_error("Cannot open input file");
    }

    fin.seekg(0, fin.end);
    size_t fileSize = fin.tellg();
    fin.seekg(0, fin.beg);
    
    vector<char> data(fileSize + 1);
    fin.read(data.data(), fileSize);
    fin.close();
    data[fileSize] = '\0';

    auto end_read = Clock::now();
    Duration read_time = end_read - start_read;
    cout << "🕐 Reading: " << read_time.count() << " seconds" << endl;

    // Parse the data
    auto start_parse = Clock::now();
    parseFromMemory(data.data());
    auto end_parse = Clock::now();
    Duration parse_time = end_parse - start_parse;
    cout << "🕑 Parsing: " << parse_time.count() << " seconds" << endl;
}

void SATLAB::parse(const char* filename) {
    cout << "Parsing file: " << filename << endl;
    ifstream file(filename, ios::binary);
    if (!file) {
        throw runtime_error("Cannot open input file");
    }

    unsigned char magic[6];
    file.read(reinterpret_cast<char*>(magic), sizeof(magic));
    file.close();

    const unsigned char xzMagic[] = {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00};

    if (memcmp(magic, xzMagic, sizeof(xzMagic)) == 0) {
        parseCompressed(filename);
    } else {
        parsePlainText(filename);
    }
}

void SATLAB::saveScopeArray(const std::string& filename, const std::vector<std::vector<Pixel>>& scope) {
    std::ofstream out(filename, std::ios::binary);
    if (!out) throw std::runtime_error("Cannot open file for writing: " + filename);
    for (const auto& row : scope) {
        out.write(reinterpret_cast<const char*>(row.data()), row.size() * sizeof(Pixel));
    }
}

void SATLAB::loadScopeArray(const std::string& filename, std::vector<std::vector<Pixel>>& scope) {
    std::ifstream in(filename, std::ios::binary);
    if (!in) throw std::runtime_error("Cannot open file for reading: " + filename);
    for (auto& row : scope) {
        in.read(reinterpret_cast<char*>(row.data()), row.size() * sizeof(Pixel));
    }
}

void SATLAB::createByteImages(const std::string& filename, cv::Mat& directImg, cv::Mat& minImg, cv::Mat& maxImg, cv::Mat& meanImg) {
    constexpr int outSize = 256;
    constexpr size_t totalBytes = outSize * outSize;
    std::vector<uint8_t> data(totalBytes, 0);
    std::vector<uint8_t> allData; // for pooled images

    // Check for xz magic header
    std::ifstream file(filename, std::ios::binary);
    if (!file) throw std::runtime_error("Cannot open file: " + filename);
    unsigned char magic[6] = {0};
    file.read(reinterpret_cast<char*>(magic), sizeof(magic));
    file.seekg(0, std::ios::beg);
    const unsigned char xzMagic[] = {0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00};
    bool is_xz = (file.gcount() == 6 && memcmp(magic, xzMagic, sizeof(xzMagic)) == 0);
    file.close();

    if (is_xz) {
        // Decompress the whole file
        FILE *inFile = fopen(filename.c_str(), "rb");
        if (!inFile) throw std::runtime_error("Cannot open input file");
        lzma_stream strm = LZMA_STREAM_INIT;
        lzma_ret ret = lzma_stream_decoder(&strm, UINT64_MAX, LZMA_CONCATENATED);
        if (ret != LZMA_OK) {
            fclose(inFile);
            throw std::runtime_error("Failed to initialize LZMA decoder");
        }
        const size_t BUFFER_SIZE = 1024 * 1024;
        std::vector<uint8_t> inBuf(BUFFER_SIZE);
        std::vector<uint8_t> outBuf(BUFFER_SIZE);
        allData.clear();
        strm.next_out = outBuf.data();
        strm.avail_out = outBuf.size();
        while (true) {
            if (strm.avail_in == 0) {
                strm.next_in = inBuf.data();
                strm.avail_in = fread(inBuf.data(), 1, inBuf.size(), inFile);
            }
            ret = lzma_code(&strm, strm.avail_in == 0 ? LZMA_FINISH : LZMA_RUN);
            size_t produced = outBuf.size() - strm.avail_out;
            if (produced > 0) {
                allData.insert(allData.end(), outBuf.data(), outBuf.data() + produced);
            }
            strm.next_out = outBuf.data();
            strm.avail_out = outBuf.size();
            if (ret == LZMA_STREAM_END) break;
            if (ret != LZMA_OK && ret != LZMA_STREAM_END) {
                lzma_end(&strm);
                fclose(inFile);
                throw std::runtime_error("Error during decompression");
            }
        }
        lzma_end(&strm);
        fclose(inFile);
        // Fill data for direct image
        size_t copyBytes = std::min(totalBytes, allData.size());
        std::copy(allData.begin(), allData.begin() + copyBytes, data.begin());
    } else {
        // Uncompressed file
        std::ifstream fin(filename, std::ios::binary);
        if (!fin) throw std::runtime_error("Cannot open file: " + filename);
        fin.seekg(0, fin.end);
        size_t n = fin.tellg();
        fin.seekg(0, fin.beg);
        allData.resize(n);
        fin.read(reinterpret_cast<char*>(allData.data()), n);
        fin.close();
        size_t copyBytes = std::min(totalBytes, allData.size());
        std::copy(allData.begin(), allData.begin() + copyBytes, data.begin());
    }
    // Create direct image
    cv::Mat directRaw;
    if (allData.size() < totalBytes) {
        // Not enough data, stretch to 256x256
        if (!allData.empty()) {
            int k = static_cast<int>(std::sqrt(allData.size()));
            if (k * k > allData.size()) --k;
            if (k == 0) throw std::runtime_error("File too small for image: " + filename);
            cv::Mat smallImg(k, k, CV_8UC1);
            std::memcpy(smallImg.data, allData.data(), k * k);
            cv::resize(smallImg, directRaw, cv::Size(outSize, outSize), 0, 0, cv::INTER_LINEAR);
        } else {
            directRaw = cv::Mat::zeros(outSize, outSize, CV_8UC1);
        }
        cv::normalize(directRaw, directRaw, 0, 255, cv::NORM_MINMAX);
        cv::cvtColor(directRaw, directImg, cv::COLOR_GRAY2BGR);
        // For small files, min/max/mean are identical to direct
        minImg = directImg.clone();
        maxImg = directImg.clone();
        meanImg = directImg.clone();
        return;
    }
    // Normal case: enough data for pooling
    directImg = cv::Mat(outSize, outSize, CV_8UC1, data.data()).clone();
    cv::normalize(directImg, directImg, 0, 255, cv::NORM_MINMAX);
    cv::cvtColor(directImg, directImg, cv::COLOR_GRAY2BGR);
    // Now pooled images from allData
    size_t n = allData.size();
    size_t k = static_cast<size_t>(std::sqrt(n));
    while (k * k > n) --k;
    if (k == 0) throw std::runtime_error("File too small for image: " + filename);
    cv::Mat rawImg(k, k, CV_8UC1);
    std::memcpy(rawImg.data, allData.data(), k * k);
    int block = k / outSize;
    if (block < 1) block = 1;
    minImg = cv::Mat(outSize, outSize, CV_8UC1);
    maxImg = cv::Mat(outSize, outSize, CV_8UC1);
    meanImg = cv::Mat(outSize, outSize, CV_8UC1);
    for (int i = 0; i < outSize; ++i) {
        for (int j = 0; j < outSize; ++j) {
            int row0 = i * block;
            int col0 = j * block;
            int row1 = std::min(row0 + block, (int)k);
            int col1 = std::min(col0 + block, (int)k);
            uint8_t minVal = 255, maxVal = 0;
            int sum = 0, count = 0;
            for (int r = row0; r < row1; ++r) {
                for (int c = col0; c < col1; ++c) {
                    uint8_t v = rawImg.at<uint8_t>(r, c);
                    if (v < minVal) minVal = v;
                    if (v > maxVal) maxVal = v;
                    sum += v;
                    ++count;
                }
            }
            minImg.at<uint8_t>(i, j) = minVal;
            maxImg.at<uint8_t>(i, j) = maxVal;
            meanImg.at<uint8_t>(i, j) = static_cast<uint8_t>(sum / std::max(1, count));
        }
    }
    cv::normalize(minImg, minImg, 0, 255, cv::NORM_MINMAX);
    cv::normalize(maxImg, maxImg, 0, 255, cv::NORM_MINMAX);
    cv::normalize(meanImg, meanImg, 0, 255, cv::NORM_MINMAX);
    cv::cvtColor(minImg, minImg, cv::COLOR_GRAY2BGR);
    cv::cvtColor(maxImg, maxImg, cv::COLOR_GRAY2BGR);
    cv::cvtColor(meanImg, meanImg, cv::COLOR_GRAY2BGR);
}
